package com.pai4j.aigc.llm.service;

import com.pai4j.aigc.llm.entity.LLMUsageEntity;
import com.pai4j.aigc.llm.repository.dao.ILLMUsageDAO;
import com.pai4j.domain.vo.llm.Usage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;

@Service
public class LLMUsageService {

    @Autowired
    private ILLMUsageDAO usageDAO;

    public void recordSuccess(String tenantId,
                              String requestId,
                              String userId,
                              Long modelId,
                              String provider,
                              String modelCode,
                              Usage usage,
                              long latencyMs,
                              Integer costCents) {
        LLMUsageEntity e = new LLMUsageEntity();
        e.setTenantId(tenantId);
        e.setRequestId(requestId);
        e.setUserId(userId);
        e.setModelId(modelId);
        e.setProvider(provider);
        e.setModelCode(modelCode);
        if (usage != null) {
            e.setTokensPrompt(usage.getPromptTokens());
            e.setTokensCompletion(usage.getCompletionTokens());
            e.setTokensTotal(usage.getTotalTokens());
        }
        e.setLatencyMs((int) Math.min(latencyMs, Integer.MAX_VALUE));
        e.setCostCents(costCents);
        e.setSuccess(1);
        e.setCreateDate(LocalDateTime.now());
        usageDAO.save(e);
    }

    public void recordSuccess(String tenantId,
                              String requestId,
                              String userId,
                              Long modelId,
                              String provider,
                              String modelCode,
                              Integer tokensPrompt,
                              Integer tokensCompletion,
                              Integer tokensTotal,
                              long latencyMs,
                              Integer costCents) {
        LLMUsageEntity e = new LLMUsageEntity();
        e.setTenantId(tenantId);
        e.setRequestId(requestId);
        e.setUserId(userId);
        e.setModelId(modelId);
        e.setProvider(provider);
        e.setModelCode(modelCode);
        e.setTokensPrompt(tokensPrompt);
        e.setTokensCompletion(tokensCompletion);
        e.setTokensTotal(tokensTotal);
        e.setLatencyMs((int) Math.min(latencyMs, Integer.MAX_VALUE));
        e.setCostCents(costCents);
        e.setSuccess(1);
        e.setCreateDate(LocalDateTime.now());
        usageDAO.save(e);
    }

    public void recordFailure(String tenantId,
                              String requestId,
                              String userId,
                              Long modelId,
                              String provider,
                              String modelCode,
                              String errorCode,
                              String errorMessage,
                              long latencyMs) {
        LLMUsageEntity e = new LLMUsageEntity();
        e.setTenantId(tenantId);
        e.setRequestId(requestId);
        e.setUserId(userId);
        e.setModelId(modelId);
        e.setProvider(provider);
        e.setModelCode(modelCode);
        e.setLatencyMs((int) Math.min(latencyMs, Integer.MAX_VALUE));
        e.setSuccess(0);
        e.setErrorCode(errorCode);
        e.setErrorMessage(errorMessage);
        e.setCreateDate(LocalDateTime.now());
        usageDAO.save(e);
    }
}
